#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

#include "libmov/fmp4-writer.h"
#include "libmov/mov-format.h"
#include "libmov/mov-reader.h"
#include "libflv/mpeg4-hevc.h"
#include "libflv/mpeg4-avc.h"

#include "plug_fmp4muxer.h"
#include "hi_ext_util.h"

#define FMP4MUXER_TASK_MAX_COUNT    (2)

typedef struct tagFMP4MUXER_VIDEO_ATTR_S
{
    FMP4MUXER_VIDEO_TYPE_E enType;
    HI_S32 s32Width;
    HI_S32 s32Height;
    union
    {
        struct mpeg4_avc_t avc;
        struct mpeg4_hevc_t hevc;
    } uPriv;
} FMP4MUXER_VIDEO_ATTR_S;

typedef struct tagFMP4MUXER_TASK_CONTEXT_S
{
    HI_BOOL bCreate;
    FILE* fp;
    fmp4_writer_t* fmp4;
    HI_S32 s32AudioTrack;
    HI_S32 s32VideoTrack;
    HI_BOOL bFirstKeyFrame;
    FMP4MUXER_VIDEO_ATTR_S stVideoAttr;
} FMP4MUXER_TASK_CONTEXT_S;

typedef struct tagFMP4MUXER_CONTEXT_S
{
    HI_BOOL bInit;
    FMP4MUXER_TASK_CONTEXT_S stTaskCtx[FMP4MUXER_TASK_MAX_COUNT];
} FMP4MUXER_CONTEXT_S;

static FMP4MUXER_CONTEXT_S s_stFMP4MuxerCtx = {
    .bInit = HI_FALSE,
};

const struct mov_buffer_t* mov_file_buffer(void);

static uint8_t s_extra_data[64 * 1024];

HI_S32 FMP4MUXER_Init(HI_VOID)
{
    if(s_stFMP4MuxerCtx.bInit == HI_TRUE)
    {
        LOGI("fmp4muxer already inited.\n");
        return HI_SUCCESS;
    }

    HI_S32 i = 0;

    for(i = 0; i < FMP4MUXER_TASK_MAX_COUNT; i++)
    {
        s_stFMP4MuxerCtx.stTaskCtx[i].bCreate = HI_FALSE;
        s_stFMP4MuxerCtx.stTaskCtx[i].fmp4 = HI_NULL;
        s_stFMP4MuxerCtx.stTaskCtx[i].fp = HI_NULL;
        s_stFMP4MuxerCtx.stTaskCtx[i].s32AudioTrack = -1;
        s_stFMP4MuxerCtx.stTaskCtx[i].s32VideoTrack = -1;
        s_stFMP4MuxerCtx.stTaskCtx[i].bFirstKeyFrame = HI_FALSE;
        memset((void *)&(s_stFMP4MuxerCtx.stTaskCtx[i].stVideoAttr), 0x0, sizeof(FMP4MUXER_VIDEO_ATTR_S));
    }

    s_stFMP4MuxerCtx.bInit = HI_TRUE;

    return HI_SUCCESS;
}

HI_S32 FMP4MUXER_DeInit(HI_VOID)
{
    if(s_stFMP4MuxerCtx.bInit == HI_FALSE)
    {
        LOGE("fmp4muxer is not init.\n");
        return HI_FAILURE;
    }

    HI_S32 i = 0;
    for(i = 0; i < FMP4MUXER_TASK_MAX_COUNT; i++)
    {
        if(s_stFMP4MuxerCtx.stTaskCtx[i].bCreate == HI_TRUE)
        {
            if(s_stFMP4MuxerCtx.stTaskCtx[i].fmp4 != HI_NULL)
            {
                fmp4_writer_destroy(s_stFMP4MuxerCtx.stTaskCtx[i].fmp4);
                s_stFMP4MuxerCtx.stTaskCtx[i].fmp4 = HI_NULL;
            }

            if(s_stFMP4MuxerCtx.stTaskCtx[i].fp != HI_NULL)
            {
                fclose(s_stFMP4MuxerCtx.stTaskCtx[i].fp);
                s_stFMP4MuxerCtx.stTaskCtx[i].fp = HI_NULL;
            }

            s_stFMP4MuxerCtx.stTaskCtx[i].bCreate = HI_FALSE;
        }
    }

    s_stFMP4MuxerCtx.bInit = HI_FALSE;

    return HI_SUCCESS;
}

HI_S32 FMP4MUXER_Create(const HI_CHAR* filename, FMP4MUXER_VIDEO_ENCODER_ATTR_S *pstVideoAttr,
                    FMP4MUXER_AUDIO_ENCODER_ATTR_S *pstAudioAttr, HI_HANDLE *hFMP4Muxer)
{
    if(s_stFMP4MuxerCtx.bInit == HI_FALSE)
    {
        LOGE("fmp4muxer is not init.\n");
        return HI_FAILURE;
    }
    
    FILE* wfp;
    fmp4_writer_t* fmp4;
    HI_S32 i = 0;

    for(i = 0; i < FMP4MUXER_TASK_MAX_COUNT; i++)
    {
        if(s_stFMP4MuxerCtx.stTaskCtx[i].bCreate != HI_TRUE)
            break;
    }

    if(i >= FMP4MUXER_TASK_MAX_COUNT)
    {
        LOGE("not available fmp4muxer task.\n");
        return HI_FAILURE;
    }

    wfp = fopen(filename, "wb");
    if(wfp == HI_NULL)
    {
        LOGE("open %s failed.\n", filename);
        return HI_FAILURE;
    }
    
    fmp4 = fmp4_writer_create(mov_file_buffer(), wfp, 0);
    if(fmp4 == HI_NULL)
    {
        LOGE("fmp4_writer_create failed.\n");
        fclose(wfp);
        return HI_FAILURE;
    }

    if(s_stFMP4MuxerCtx.stTaskCtx[i].s32AudioTrack == -1)
    {
        s_stFMP4MuxerCtx.stTaskCtx[i].s32AudioTrack = fmp4_writer_add_audio(fmp4,
                        MOV_OBJECT_AAC, 
                        pstAudioAttr->s32ChannelCount,
                        pstAudioAttr->s32BitsPerSample,
                        pstAudioAttr->s32SampleRate,
                        HI_NULL, 0);
        if(s_stFMP4MuxerCtx.stTaskCtx[i].s32AudioTrack < 0)
        {
            LOGE("fmp4_writer_add_audio failed.\n");
            s_stFMP4MuxerCtx.stTaskCtx[i].s32AudioTrack = -1;
            fmp4_writer_destroy(fmp4);
            fclose(wfp);
            return HI_FAILURE;
        }
    }

    LOGD("audio track id : %d\n", s_stFMP4MuxerCtx.stTaskCtx[i].s32AudioTrack);

    s_stFMP4MuxerCtx.stTaskCtx[i].stVideoAttr.enType = pstVideoAttr->enType;
    s_stFMP4MuxerCtx.stTaskCtx[i].stVideoAttr.s32Width = pstVideoAttr->s32Width;
    s_stFMP4MuxerCtx.stTaskCtx[i].stVideoAttr.s32Height = pstVideoAttr->s32Height;

    s_stFMP4MuxerCtx.stTaskCtx[i].fp = wfp;
    s_stFMP4MuxerCtx.stTaskCtx[i].fmp4 = fmp4;
    s_stFMP4MuxerCtx.stTaskCtx[i].bCreate = HI_TRUE;
    
    *hFMP4Muxer = i;
    return HI_SUCCESS;
}

HI_S32 FMP4MUXER_Close(HI_HANDLE hFMP4Muxer)
{
    if(s_stFMP4MuxerCtx.bInit == HI_FALSE)
    {
        LOGE("fmp4muxer is not init.\n");
        return HI_FAILURE;
    }

    if(hFMP4Muxer < 0 || hFMP4Muxer >= FMP4MUXER_TASK_MAX_COUNT)
    {
        LOGE("invalid fmp4muxer %u\n", hFMP4Muxer);
        return HI_FAILURE;
    }

    if(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].bCreate != HI_TRUE)
    {
        LOGE("task %u not created.\n", hFMP4Muxer);
        return HI_FAILURE;
    }

    if(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].fmp4 != HI_NULL)
    {
        fmp4_writer_destroy(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].fmp4);
        s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].fmp4 = HI_NULL;
    }

    if(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].fp != HI_NULL)
    {
        fclose(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].fp);
        s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].fp = HI_NULL;
    }

    memset((void *)&(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].stVideoAttr), 0x0, sizeof(FMP4MUXER_VIDEO_ATTR_S));
    s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].stVideoAttr.enType = FMP4MUXER_VIDEO_TYPE_BUTT;
    s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].s32AudioTrack = -1;
    s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].s32VideoTrack = -1;
    s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].bFirstKeyFrame = HI_FALSE;
    s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].bCreate = HI_FALSE;

    return HI_SUCCESS;
}

HI_S32 FMP4MUXER_Write(HI_HANDLE hFMP4Muxer, FMP4MUXER_TRACK_TYPE_E enTrackType,
                    const HI_VOID* data, HI_U32 bytes, HI_S64 pts, HI_S64 dts, HI_S32 flags)
{
    if(s_stFMP4MuxerCtx.bInit == HI_FALSE)
    {
        LOGE("fmp4muxer is not init.\n");
        return HI_FAILURE;
    }
    
    if(hFMP4Muxer < 0 || hFMP4Muxer >= FMP4MUXER_TASK_MAX_COUNT)
    {
        LOGE("invalid fmp4muxer %u\n", hFMP4Muxer);
        return HI_FAILURE;
    }

    if(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer].bCreate != HI_TRUE)
    {
        LOGE("task %u not created.\n", hFMP4Muxer);
        return HI_FAILURE;
    }

    HI_S32 s32Ret = HI_SUCCESS, n = 0;
    FMP4MUXER_TASK_CONTEXT_S *pFMP4MuxerTaskCtx = &(s_stFMP4MuxerCtx.stTaskCtx[hFMP4Muxer]);

    if(enTrackType == FMP4MUXER_TRACK_TYPE_VIDEO)
    {
        HI_S32 vcl = 0, update = 0;
        HI_U8 nalutype = 0;
        HI_CHAR *output;

        output = malloc(bytes);
        if(output == HI_NULL)
        {
            LOGE("malloc failed.\n");
            return HI_FAILURE;
        }

        if(pFMP4MuxerTaskCtx->stVideoAttr.enType == FMP4MUXER_VIDEO_TYPE_H264)
        {
            struct mpeg4_avc_t *pAvc;

            pAvc = &(pFMP4MuxerTaskCtx->stVideoAttr.uPriv.avc);

            n = h264_annexbtomp4(pAvc, data, bytes, output, bytes, &vcl, &update);
            if(n <= 0)
            {
                LOGE("h264_annexbtomp4 failed.\n");
                free(output);
                return HI_FAILURE;
            }

            if(pFMP4MuxerTaskCtx->s32VideoTrack < 0)
            {
                if (pAvc->nb_sps < 1 || pAvc->nb_pps < 1)
                {
                    free(output);
                    return -2; // waiting for sps/pps
                }

                int extra_data_size = mpeg4_avc_decoder_configuration_record_save(pAvc, s_extra_data, sizeof(s_extra_data));
                if (extra_data_size <= 0)
                {
                    // invalid AVCC
                    free(output);
                    return HI_FAILURE;
                }

                pFMP4MuxerTaskCtx->s32VideoTrack = fmp4_writer_add_video(pFMP4MuxerTaskCtx->fmp4,
                                                            MOV_OBJECT_H264,
                                                            pFMP4MuxerTaskCtx->stVideoAttr.s32Width,
                                                            pFMP4MuxerTaskCtx->stVideoAttr.s32Height,
                                                            s_extra_data, extra_data_size);
                if (pFMP4MuxerTaskCtx->s32VideoTrack < 0)
                {
                    free(output);
                    return HI_FAILURE;
                }

                LOGD("video track id : %d\n", pFMP4MuxerTaskCtx->s32VideoTrack);
            }

            nalutype = ((HI_U8 *)data)[4] & 0x1f;

            if(pFMP4MuxerTaskCtx->bFirstKeyFrame == HI_FALSE && nalutype == 5)
            {
                pFMP4MuxerTaskCtx->bFirstKeyFrame = HI_TRUE;
            }

            if(nalutype >= 1 && nalutype <= 5)
            {
                s32Ret = fmp4_writer_write(pFMP4MuxerTaskCtx->fmp4,
                                           pFMP4MuxerTaskCtx->s32VideoTrack,
                                           output, n, pts,
                                           dts, 1 == vcl ? MOV_AV_FLAG_KEYFREAME : 0);
                if(s32Ret != HI_SUCCESS)
                {
                    LOGE("fmp4_writer_write %d\n", s32Ret);
                    return HI_FAILURE;
                }
            }
        }
        else
        {
            struct mpeg4_hevc_t *pHevc;

            pHevc = &(pFMP4MuxerTaskCtx->stVideoAttr.uPriv.hevc);
            
            n = h265_annexbtomp4(pHevc, data, bytes, output, bytes, &vcl, &update);
            if(n <= 0)
            {
                LOGE("h265_annexbtomp4 failed.\n");
                free(output);
                return HI_FAILURE;
            }

            if(pFMP4MuxerTaskCtx->s32VideoTrack < 0)
            {
                if (pHevc->numOfArrays < 3)
                {
                    free(output);
                    return -2; // waiting for vps/sps/pps
                }

                int extra_data_size = mpeg4_hevc_decoder_configuration_record_save(pHevc, s_extra_data, sizeof(s_extra_data));
                if (extra_data_size <= 0)
                {
                    // invalid HVCC
                    free(output);
                    return HI_FAILURE;
                }

                pFMP4MuxerTaskCtx->s32VideoTrack = fmp4_writer_add_video(pFMP4MuxerTaskCtx->fmp4, 
                                                            MOV_OBJECT_HEVC,
                                                            pFMP4MuxerTaskCtx->stVideoAttr.s32Width,
                                                            pFMP4MuxerTaskCtx->stVideoAttr.s32Height,
                                                            s_extra_data, extra_data_size);
                if (pFMP4MuxerTaskCtx->s32VideoTrack < 0)
                {
                    free(output);
                    return HI_FAILURE;
                }

                LOGD("video track id : %d\n", pFMP4MuxerTaskCtx->s32VideoTrack);
            }

            nalutype = (((HI_U8 *)data)[4] >> 1) & 0x3f;

            if(pFMP4MuxerTaskCtx->bFirstKeyFrame == HI_FALSE && nalutype == 19)
            {
                pFMP4MuxerTaskCtx->bFirstKeyFrame = HI_TRUE;
            }

            if(nalutype <= 31)
            {
                s32Ret = fmp4_writer_write(pFMP4MuxerTaskCtx->fmp4,
                                           pFMP4MuxerTaskCtx->s32VideoTrack,
                                           output, n, pts,
                                           dts, 1 == vcl ? MOV_AV_FLAG_KEYFREAME : 0);
                if(s32Ret != HI_SUCCESS)
                {
                    LOGE("fmp4_writer_write %d\n", s32Ret);
                    return HI_FAILURE;
                }
            }
        }

        free(output);
        output = HI_NULL;
    }
    else if(enTrackType == FMP4MUXER_TRACK_TYPE_AUDIO)
    {
        if(pFMP4MuxerTaskCtx->s32AudioTrack == -1)
        {
            LOGE("audio track not add.\n");
            return HI_FAILURE;
        }

        if(pFMP4MuxerTaskCtx->bFirstKeyFrame == HI_TRUE)
        {
            s32Ret = fmp4_writer_write(pFMP4MuxerTaskCtx->fmp4,
                                       pFMP4MuxerTaskCtx->s32AudioTrack,
                                       data, bytes, pts,
                                       dts, flags);
            if(s32Ret != HI_SUCCESS)
            {
                LOGE("fmp4_writer_write %d\n", s32Ret);
                return HI_FAILURE;
            }
        }
    }

    return HI_SUCCESS;
}

